# -*- coding: utf-8 -*-
#refs to: machine learning in action 14.6 chapter


import numpy as np


# a^2 + b^2 + c^2 + ...
def sum_of_squares(A):
    ret = 0
    for a in A:
        ret += (a * a)
    return ret

#return the number N, top N sum_of_squares will take in all sum_of_squares as rate value
def rate_of_sum_of_squares(A,rate=1.0):
    sumOfSquares = sum_of_squares(A)
    sumOfSquaresWithRate = sumOfSquares * rate
    ret = 0
    i = 1
    for a in A:
        aa = a*a
        if ret + aa >= sumOfSquaresWithRate:
            break
        ret += aa
        i += 1
    return i

def compression_by_SVD(data,recover_rate = 1.0):
    #s = np.shape(data)
    #w = s[0]
    #h = s[1]
    #c0 = data[:, :, 0]
    #c1 = data[:, :, 1]
    #c2 = data[:, :, 2]
    #c3 = data[:, :, 3]
    U, Sigma, VT = np.linalg.svd(data)

    N = rate_of_sum_of_squares(Sigma, recover_rate)
    print "top " + str(N) + " of " + str(len(Sigma)) + \
          " will recover " + str(recover_rate * 100) + "% data"

    SigmaN = np.zeros((N,N))
    for i in range(N):
        SigmaN[i][i] = Sigma[i]
    tmp = np.dot(U[:, :N], SigmaN)
    return np.dot(tmp ,VT[:N,:])


if __name__ == '__main__':
    from PIL import Image
    import matplotlib.pyplot as plt
    img=np.array(Image.open('../dataset/kk1.jpg'))  #打开图像并转化为数字矩阵
    img_shape  = img.shape
    print img_shape
    img_channels_0 = compression_by_SVD(img[:, :, 0], 0.9)
    img_channels_1 = compression_by_SVD(img[:, :, 1], 0.99)
    img_channels_2 = compression_by_SVD(img[:, :, 2], 0.999)


    plt.subplot(221)
    plt.title("channels 0 with 0.9")
    plt.imshow(img_channels_0)

    plt.subplot(222)
    plt.title("channels 1 with 0.99")
    plt.imshow(img_channels_1)

    plt.subplot(223)
    plt.title("channels 2 with 0.999")
    plt.imshow(img_channels_2)

    img_reshape = list(img_shape)
    img_reshape[2] = 1
    img_channels_0 = img_channels_0.reshape(img_reshape)
    img_channels_1 = img_channels_1.reshape(img_reshape)
    img_channels_2 = img_channels_2.reshape(img_reshape)

    img_concated = np.concatenate([img_channels_0,img_channels_1,img_channels_2],axis=2)


    #plt.figure("dog")
    plt.subplot(224)
    #plt.imshow(img_concated)
    plt.imshow(img)
    #plt.axis('off')
    plt.show()
